import os
from asyncio import Semaphore

from langchain_openai import ChatOpenAI

from data_dyne.config import config


# GLM-4-Air 默认并发100
# GLM-4-Flash 默认并发200
# GLM-4-Plus 默认并发50
# GLM-4-FlashX 默认并发50
# GLM-Zero-Preview 默认并发50
# GLM-4-Air-0111 默认并发50
# Rerank 默认并发50
# embedding3 默认并发50

semaphore_air = Semaphore(100)
semaphore_plus = Semaphore(50)
semaphore_rerank = Semaphore(20)
semaphore_rag = Semaphore(40)

model = "glm-4-plus"

os.environ["ZHIPUAI_API_KEY"] = config.zhipuai_api_key
zhipu_llm_plus = ChatOpenAI(
    temperature=0.6,
    model=model,
    openai_api_key=config.zhipuai_api_key,
    openai_api_base="https://open.bigmodel.cn/api/paas/v4/",
)

air_model = "glm-4-air"
zhipu_llm_air = ChatOpenAI(
    temperature=0.6,
    model=air_model,
    openai_api_key=config.zhipuai_api_key,
    openai_api_base="https://open.bigmodel.cn/api/paas/v4/",
)
